'''
predict.py有几个注意点
1、无法进行批量预测，如果想要批量预测，可以利用os.listdir()遍历文件夹，利用Image.open打开图片文件进行预测。
2、如果想要保存，利用r_image.save("img.jpg")即可保存。
3、如果想要原图和分割图不混合，可以把blend参数设置成False。
4、如果想根据mask获取对应的区域，可以参考detect_image中，利用预测结果绘图的部分。
seg_img = np.zeros((np.shape(pr)[0],np.shape(pr)[1],3))
for c in range(self.num_classes):
    seg_img[:, :, 0] += ((pr == c)*( self.colors[c][0] )).astype('uint8')
    seg_img[:, :, 1] += ((pr == c)*( self.colors[c][1] )).astype('uint8')
    seg_img[:, :, 2] += ((pr == c)*( self.colors[c][2] )).astype('uint8')
'''
import traceback
from progress.bar import IncrementalBar
from FPS_test import *
import cv2

python_file = os.path.dirname(__file__)
results = os.path.join(python_file, 'results')
masks_path = os.path.join(results, 'mask-original')
original = os.path.join(python_file, 'original')
results_path = os.path.join(results, 'blend-fitted')
blend_path = os.path.join(results, 'blend-original')
compare_path = os.path.join(results, 'compare')
masks = os.path.join(python_file, 'mask')

def make_dirs(dirs):
    for dir in dirs:
        if not os.path.exists(dir):
            os.makedirs(dir)
            
def Img2cv2(pic):
    return cv2.cvtColor(np.array(pic), cv2.COLOR_RGB2BGR)

def main(pic, net):
    try:
        image = Image.open(pic)
        blend, mask = net.detect_image(image)
        if not blend is None:
            blend.save(pic.replace(original, blend_path))
        mask.save(pic.replace(original, masks_path))
    except:
        print('{:*^60}'.format('ERROR'))
        print('Open Error: ' + pic)
        exstr = traceback.format_exc()
        print(exstr)
        print('{:*^60}'.format(' '))

if __name__ == '__main__':
    start = time.time()
    pics = glob.glob(os.path.join(original, '*.jpg'))
    make_dirs([results, masks_path, blend_path])
    bar = IncrementalBar(max=len(pics))
    Net = Pspnet()
    for pic in pics:
        main(pic, Net)
        bar.next()
    end = time.time()
    print('\nTIME: ', end - start)
    FPS_test()